import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import pearsonr
from scipy.stats import zscore


class FactorEngine:
    """
    因子研究管线
        1. 支持传入单个因子
        2. 计算因子IC，IR时间序列值
        3. 计算分组收益，因子收益
        4. 可视化
    param:
        factor_name: 因子列名称列表，默认为单因子 ('signal',) 需要传入元组
        factor_raw: 原始因子数据, columns=['trading', 'code', 'signal', 'close', 'open']
        n: int, 因子延迟期数，默认为1
        groupN: int, 分组数，默认为5
    """

    def __init__(self, n=1, factor_name=None, factor_raw: pd.DataFrame = pd.DataFrame(), groupN=5, rebalance_period=20,
                 price_col='close'):

        if factor_name is None:
            self.factor_name = ('signal',)
        else:
            self.factor_name = factor_name

        self.data = factor_raw  # 原始因子值
        self.data['trading'] = pd.to_datetime(self.data['trading'])
        self.data.sort_values(by=['trading'], inplace=True)
        self.data.reset_index(drop=True, inplace=True)

        self.factor_result = {'ic': {}, 'ir': {}}  # 'ic': {因子名:{数据帧时间: IC值}}

        self.n = n  # 因子延迟期数
        self.groupN = groupN
        self.holdings = {}  # 按调仓周期的因子持仓记录
        self.group_return = {}  # 按调仓周期的因子分组回测收益率
        self.rebalance_period = rebalance_period  # 分组回测调仓频率
        self.price_col = price_col  # 收盘价列名

        # 因子数据处理
        _temp = 'close_' + str(self.n)
        self.data[_temp] = self.data.groupby('code')[self.price_col].shift(-self.n)  # 获取t+n日收盘价
        self.close_shift = _temp  # 记录滞后收盘价列名
        self.data.dropna(axis=0, how='any', inplace=True)  # 剔除空行

        # copy and basic checks
        self.data = factor_raw.copy()
        if 'trading' not in self.data.columns or 'code' not in self.data.columns:
            raise ValueError("factor_raw must contain 'trading' and 'code' columns")

        self.data['trading'] = pd.to_datetime(self.data['trading']).dt.date
        self.data.sort_values(['trading', 'code'], inplace=True)
        self.data.reset_index(drop=True, inplace=True)

        if len(self.data['trading'].unique()) < 6:
            raise ValueError('因子数据长度应大于6个交易日')

        # create forward price and forward returns
        self.data[f'return_{self.n}d'] = self.data[self.close_shift] / self.data[self.price_col] - 1

        self.next_return_col = f'return_{self.n}d'

        # drop rows without forward return (end of series)
        self.data.dropna(subset=[f'return_{self.n}d'], inplace=True)

        # containers
        self.factor_result = {
            'ic': {},  # ic timeseries dict per factor
            'ir': {},  # ir per factor
            'ic_spearman': {},
            'ic_pearson': {},
            'ic_rank': {},
            'rolling_ic': {},
        }
        self.group_return = {}
        self.holdings = {}

    def prepare_factor_data(self, signal_name: str):
        """
        计算因子IC值
        :param signal_name: 因子名称
        :return:
        """

        _factor_data = self.data[['trading', 'code', signal_name, 'open', 'close', self.next_return_col]].copy()
        # _factor_data = _factor_data[~np.isinf(_factor_data[signal_name])]  # 剔除异常因子值

        # 计算IC
        _ic = {}
        _ic_mean = []

        _tradings = _factor_data['trading'].unique()
        for trade in _tradings:
            _signal = _factor_data[_factor_data['trading'] == trade][signal_name].values
            _nextReturn = _factor_data[_factor_data['trading'] == trade][self.next_return_col].values

            # 检查是否为常量序列
            if np.all(_signal == _signal[0]) or np.all(_nextReturn == _nextReturn[0]):
                _r = np.nan
            else:
                _r = pearsonr(_signal, _nextReturn)[0]  # 计算皮尔森相关系数
            _ic_mean.append(_r)
            _ic.update({trade: _r})  # 每数据帧ic

        # 计算IR，IR = mean(ic) / std(ic)
        _ir = np.nanmean(_ic_mean) / np.nanstd(_ic_mean)  # 忽略空值

        self.factor_result['ic'].update({signal_name: _ic})
        self.factor_result['ir'].update({signal_name: _ir})

    def group_report(self, signal_name):
        """
            因子分组回测
            :param signal_name: 因子名称
        """
        rebalance_period = self.rebalance_period
        # 设置调仓日
        set_trading = self.data['trading'].unique()
        _rebalance = set_trading[::rebalance_period].tolist()  # 按调仓日期间隔采样
        if set_trading[-1] not in _rebalance:
            _rebalance.append(set_trading[-1])

        # 交易日预分组
        trading_groups = {date: group for date, group in self.data.groupby('trading')}

        # 按调仓日进行遍历计算
        _result = {}
        _port = []

        for date0, date1 in zip(_rebalance[:-1], _rebalance[1:]):
            # 调仓日对因子值进行分组
            _data0 = trading_groups[date0].copy()
            _data1 = trading_groups[date1].copy()

            factor_series = _data0[signal_name]
            valid_data = factor_series.dropna()

            try:
                _data0.loc[valid_data.index, 'group'] = pd.qcut(
                    valid_data, self.groupN, labels=np.arange(1, self.groupN + 1),
                    duplicates='drop'
                )
            except Exception as e:
                print(f"分组异常，有效样本数{len(_data0)}, 唯一值数据{_data0[signal_name].unique()}, 日期：{date0}")
                print(f"报错：{e}")
                try:
                    ranked = valid_data.rank(pct=True)
                    bins = pd.cut(ranked, self.groupN, labels=np.arange(1, self.groupN + 1))
                    _data0.loc[valid_data.index, 'group'] = bins
                    print(f"{date0}, 采用rank+cut方式进行分组")
                except Exception as e:
                    raise ValueError(f"依旧分组失败，报错{e}")

            _port.append(_data0)
            # 计算每组下一期收益率
            _group_return = []  # 记录顺序为group1, group2, group3, ...

            for group in range(1, self.groupN + 1):
                _code = _data0[_data0['group'] == group]['code'].to_list()  # 组内股票列表
                _group_data0 = _data0[_data0['group'] == group].copy()
                _group_data1 = _data1[_data1['code'].isin(_code)].copy()
                _group = pd.concat([_group_data0, _group_data1], ignore_index=True)
                _group.sort_values(by='trading', inplace=True)
                _return = _group.groupby(['code'], observed=True)[self.price_col].apply(
                    lambda x: x.iloc[-1] / x.iloc[0] - 1).mean()  # 分组的下一期收益率
                _group_return.append(_return)

            _result.update({date1: _group_return})
        _port = pd.concat(_port, ignore_index=True)

        self.group_return.update({signal_name: _result})
        self.holdings.update({signal_name: _port})

    def run(self, signal_name='signal'):
        self.prepare_factor_data(signal_name, )
        self.group_report(signal_name, )

    def plotting_group_return(self, factor_name="signal"):
        """输出分组累计收益率曲线"""
        # 获取分组收益率数据
        _data = pd.DataFrame.from_dict(self.group_return[factor_name], orient='index',
                                       columns=np.arange(1, self.groupN + 1))

        fig, ax = plt.subplots(figsize=(6, 4))
        for group in range(1, self.groupN + 1):
            _group = (1 + _data[group]).cumprod() - 1
            # 添加曲线
            ax.plot(_data.index, _group, label=f"group_{group}")

        ax.set_title(f"{factor_name}: accumulated_return_grouped, rebalance_period={self.rebalance_period}")
        plt.figtext(0.25, 0.01,
                    "Group according to factor values from small to large")
        ax.set_xlabel("trade_date")
        ax.set_ylabel("cumulative return")
        ax.legend(loc='best')
        plt.show()

    def plotting_ic(self, factor_name="signal", rolling_window=20):
        """输出IC移动平均曲线，相关统计指标"""
        # 获取ic，ir数据
        _ic = pd.DataFrame.from_dict(self.factor_result['ic'][factor_name], orient='index', columns=['ic'])
        _ic.loc[:, 'ic_ma'] = _ic['ic'].rolling(window=rolling_window).mean()
        _ir = self.factor_result['ir'][factor_name]

        # 设置画布
        fig, axes = plt.subplots(1, 2, figsize=(6, 4))

        # ic时间序列
        axes[0].plot(_ic.index, _ic['ic'], label='timeseries_ic')
        axes[0].set_title(f'{factor_name}: timeseries_ic')
        axes[0].set_xlabel('trade_date')
        axes[0].set_ylabel('IC')

        # ic移动平均线
        axes[1].plot(_ic.index[rolling_window - 1:], _ic['ic_ma'][rolling_window - 1:], label=f'ma{rolling_window}_ic')
        axes[1].set_title(f'{factor_name}: ma_ic')
        axes[1].set_xlabel('trade_date')
        axes[1].set_ylabel('IC_ma')

        # 底部打印统计指标

        plt.figtext(0.5, 0.01,
                    f"Mean of IC: {_ic['ic'].mean(): .4f}, Positive proportion: {len(_ic[_ic['ic'] > 0]) / len(_ic): .4f}\n"
                    f"IR: {_ir: .4f}",
                    ha='center', fontsize=10, bbox={"facecolor": "orange", "alpha": 0.5, "pad": 5})
        plt.tight_layout()
        plt.show()


def factor_correlation(factors_ic):
    """
    计算因子间的相关系数
    实现方法：
        因子IC序列之间的spearman相关系数
    ：param factors_ic: 因子IC序列，dataframe.columns=[trading, factor1, factor2, ...]
    :return: 相关系数矩阵，
    """
    _factors = factors_ic.columns[1:].to_list()
    _result = {'factor_name': _factors}
    for i, factor in enumerate(_factors):
        # 每一个因子都与所有因子计算相关性
        _f1 = factors_ic[factor].values
        _r = []
        for j in _factors:
            _f2 = factors_ic[j].values
            _corr = pearsonr(_f1, _f2)[0]
            _r.append(_corr)
        _result[factor] = _r
    return pd.DataFrame(data=_result)


def factor_return(df):
    """
        因子收益率计算
        - 原理：下期收益率与因子值（标准化后）截面回归，回归系数为因子收益率
        - 向量化实现：最小二乘目标小的数值解，因子暴露与下期收益的内积/因子暴露二范数的平方
            - 因子暴露矩阵X=(T, N)

        Input:
        * df: columns=[trade_date, code, signal, daily_return]
    """
    required_columns = {'trade_date', 'code', 'signal', 'daily_return'}

    if not required_columns.issubset(df.columns):
        missing_cols = required_columns - set(df.columns)
        raise ValueError(f"缺失必要列: {missing_cols}")

    # 因子暴露矩阵
    df.sort_values(by='trade_date', inplace=True, ascending=True)
    _date = df.trade_date.unique()

    def _std(group):
        # nan进行0填充
        group['signal'] = group['signal'].fillna(0)
        group['std_signal'] = zscore(group['signal'])
        return group

    def _fillna(a):
        _mean = a.mean()
        a.fillna(_mean, inplace=True)
        return a

    df1 = df.groupby('trade_date', as_index=False).apply(lambda group: _std(group))
    df1.reset_index(inplace=True, drop=True)

    # 构建因子暴露矩阵, 每一行处理nan用均值替代
    _x = df1.pivot(index='trade_date', columns='code', values='std_signal').iloc[:-1, :]
    _x = _x.apply(lambda a: _fillna(a), axis=1).values

    # 构建下期个股收益矩阵
    _y = df1.pivot(index='trade_date', columns='code', values='daily_return').iloc[1:, :]
    _y = _y.apply(lambda a: _fillna(a), axis=1).values

    _numerator = np.nansum(_x * _y, axis=1)
    _denominator = np.nansum(_x ** 2, axis=1)

    _factor_return = _numerator / _denominator

    return pd.Series(data=_factor_return, index=_date[1:])
